import copy
import pandas as pd
import numpy as np
from statsmodels import api as sm
from matplotlib import pyplot as plt


def de_trend(ts, method='hp', ts_freq='M', if_plot=False, **kwargs):
    """
    序列进行去趋势

    :param pd.Series ts: 时间序列指标
    :param str method: 去趋势方法，支持'hp'， 'ma'; 均线法是以均线作为趋势项
    :param str ts_freq: 时间序列频率， 支持'Y', 'Q', 'M', 'D'; 默认为'M'
    :param bool if_plot: 是否画图
    :param kwargs:
        - str file_name: 输出图片名
        - float lamda:
        - int periods: 均线滤波对应的参数
        - int min_periods: 均线滤波对应的参数
    :return:
    """

    if method == 'hp':
        lamda = kwargs.pop('lamda', None)
        if lamda is None:
            if ts_freq == 'Q':
                lamda = 1600
            elif ts_freq == 'Y':
                lamda = 1600 * (1 / 4) ** 4
            elif ts_freq == 'M':
                lamda = 1600 * 3 ** 4
            elif ts_freq == 'D':
                lamda = 1600 * 9 ** 4
        cycle, trend = sm.tsa.filters.hpfilter(ts, lamda)
        cycle.name = ts.name
        trend.name = ts.name
    elif method == 'ma':
        period = kwargs.pop('periods', 120)
        min_periods = kwargs.pop('min_periods', period)
        trend = ts.rolling(period, min_periods).mean()
        cycle = ts - trend
    else:
        raise NotImplementedError

    if if_plot:
        f, ax = plt.subplots(figsize=(15, 3))
        ax.plot(ts, label='signal')
        ax.plot(trend, label='trend')
        ax.plot(cycle, label='cycle')
        ax.set(title=f'de_trend')
        plt.legend()

        file_name = kwargs.pop('file_name', None)
        if file_name is not None:
            plt.savefig(f'{file_name}.png', dpi=400, bbox_inches='tight')

        plt.show()

    return cycle, trend


def de_noise(ts, method='hp', ts_freq='M', if_plot=False, **kwargs):
    """
    去除指标序列中的高频噪音

    :param pd.Series ts: 指标值序列
    :param str method: 去噪音方法， 支持''hp', 'ma', 'ema', 'mm', 'llt'
    :param str ts_freq: 数据频率，支持'Y'， 'Q'， 'M'， 'D'
    :param bool if_plot: 是否画图，默认为False, 不画图
    :param kwargs:
        - str file_name: 图片保存名字, 默认为None, 即不保存图片
        - int lamda: hp滤波法需要用到的参数
        - int periods: 均线滤波对应的参数
        - int min_periods: 均线滤波对应的参数
        - int d: llt滤波对应的参数, 默认为12
    :return:
    """

    if method == 'hp':
        lamda = kwargs.pop('lamda', None)
        if lamda is None:
            if ts_freq == 'Q':
                lamda = 1
            elif ts_freq == 'Y':
                lamda = 1 * (1 / 4) ** 4
            elif ts_freq == 'M':
                lamda = 1 * 3 ** 4
            elif ts_freq == 'D':
                lamda = 1 * 9 ** 4
        noise, denoised_signal = sm.tsa.filters.hpfilter(ts, lamda)
        denoised_signal.name = ts.name
    elif method == 'ma':
        periods = kwargs.pop('periods', 7)
        min_periods = kwargs.pop('min_periods', periods)
        denoised_signal = ts.rolling(periods, min_periods).mean()
        noise = ts - denoised_signal
    elif method == 'ema':
        periods = kwargs.pop('periods', 7)
        min_periods = kwargs.pop('min_periods', periods)
        denoised_signal = ts.ewm(span=periods, min_periods=min_periods, adjust=True).mean()
        noise = ts - denoised_signal
    elif method == 'mm':
        periods = kwargs.pop('periods', 7)
        min_periods = kwargs.pop('min_periods', periods)
        denoised_signal = ts.rolling(periods, min_periods).median()
        noise = ts - denoised_signal
    elif method == 'llt':
        d = kwargs.pop('d', 12)
        cp = copy.deepcopy(ts)
        tmp = copy.deepcopy(ts)
        a = 2 / (d + 1)
        for i in range(2, len(tmp)):
            tmp[i] = (a - (a * a) / 4) * cp[i] + ((a * a) / 2) * cp[i - 1] - (a - 3 * (a * a) / 4) * cp[
                i - 2] + 2 * (1 - a) * tmp[i - 1] - (1 - a) * (1 - a) * tmp[i - 2]
            denoised_signal = tmp
            noise = ts - denoised_signal
    else:
        raise NotImplementedError

    if if_plot:
        f, ax = plt.subplots(figsize=(15, 3))
        ax.plot(ts, label='signal')
        ax.plot(denoised_signal, label='denoised_signal')
        ax.plot(noise, label='noise')
        ax.set(title=f'de_noise')
        plt.legend()

        file_name = kwargs.pop('file_name', None)
        if file_name is not None:
            plt.savefig(f'{file_name}.png', dpi=400, bbox_inches='tight')

        plt.show()

    return denoised_signal, noise


def ts_percentile_rank_score(ts, way='expanding', rank_method='rank', periods=None, min_periods=1, positive=True,
                             scale=10, preprocess=False, **kwargs):
    """

    :param pd.Series ts: 指标序列
    :param str way: 排名百分位计算方式， 支持‘expanding’和‘rolling’, 'whole'
                    expanding, 则依据指标全样本数据确定分位点（会利用到未来数据）
                    rolling, 则使用滚动估计的方法 不断扩展样本形成动态的分位点 避免了使用未来数据
                    whole, 全历史数据排序得分
    :param str rank_method: 排名方法，支持'rank' 即指标排序百分比，'quantile' 即指标值所处的分位点
    :param int periods: 滚动窗口长度
    :param int min_periods: 窗口最小长度
    :param bool positive: 方向 - 正序排还是逆序排
    :param int scale: 转换成几分制
    :param kwargs:
        - bool if_detrend: 默认为False
        - bool if_denoise: 默认为True， if_detrend和if_denoise这个至少一个为True
        - str detrend_method: 默认为hp
        - str denoise_method: 默认为hp
        - dict detrend_kwargs: {}
        - dict denoise kwargs: {}

    :return: 返回scale分制得分，得分越高越好
    """

    data = ts.copy()

    if_detrend = kwargs.pop('if_detrend', False)
    if_denoise = kwargs.pop('if_denoise', True)
    detrend_method = kwargs.pop('detrend_method', 'hp')
    detrend_kwargs = kwargs.pop('detrend_kwargs', {})
    denoise_method = kwargs.pop('denoise_method', 'hp')
    denoise_kwargs = kwargs.pop('denoise_kwargs', {})

    if len(data.shape) == 1:  # data为纵向
        if data.unique().shape[0] == 1:  # series中所有值都一样
            return pd.Series([np.nan for i in range(len(data))], index=data.index)
    else:  # data为横向
        if data.shape[1] == 1:
            data = data.T.squeeze()

    if if_detrend and detrend_method == 'ma':
        trend_periods = detrend_kwargs['periods'] if 'periods' in detrend_kwargs.keys() else 120
        trend_min_periods = detrend_kwargs['min_periods'] if 'min_periods' in detrend_kwargs.keys() else 120
        ma_trend = data.rolling(trend_periods, trend_min_periods).mean()

    def last_rank_score(x):
        if rank_method == 'rank':
            rank = x.rank(pct=True, ascending=positive).iloc[-1]
        elif rank_method == 'quantile':
            rank_quantile = (x - x.min()) / (x.max() - x.min())
            if positive:
                rank = rank_quantile.iloc[-1]
            else:
                rank = 1 - rank_quantile.iloc[-1]
        return rank

    def last_rank_score_with_preprocess(x):
        if if_detrend:
            if detrend_method == 'hp':
                x, x_trend = de_trend(x, method=detrend_method, ts_freq='D', **detrend_kwargs)
            elif detrend_method == 'ma':
                x = x - ma_trend
            x.dropna(inplace=True)
        if if_denoise:
            x, x_noise = de_noise(x, method=denoise_method, ts_freq='D', **denoise_kwargs)
            x.dropna(inplace=True)

        if rank_method == 'rank':
            rank = x.rank(pct=True, ascending=positive).iloc[-1]
        elif rank_method == 'quantile':
            rank_quantile = (x - x.min()) / (x.max() - x.min())
            if positive:
                rank = rank_quantile.iloc[-1]
            else:
                rank = 1 - rank_quantile.iloc[-1]

        return rank

    if preprocess:
        func = last_rank_score_with_preprocess
    else:
        func = last_rank_score

    if way == 'expanding' or way == 'rolling':
        if if_detrend and detrend_method == 'ma':
            data = data[trend_min_periods - 1:]

    if way == 'expanding':
        rank = data.expanding(min_periods).apply(func)
    elif way == 'rolling':
        rank = data.rolling(periods, min_periods).apply(func)
    elif way == 'whole':
        if preprocess:
            if if_detrend:
                data, x_trend = de_trend(data, method=detrend_method, ts_freq='D', **detrend_kwargs)
                data.dropna(inplace=True)
            if if_denoise:
                data, x_noise = de_noise(data, method=denoise_method, ts_freq='D', **denoise_kwargs)
            if rank_method == 'rank':
                rank = data.rank(pct=True, ascending=positive)
            elif rank_method == 'quantile':
                rank_quantile = (data - data.min()) / (data.max() - data.min())
                if positive:
                    rank = rank_quantile
                else:
                    rank = 1 - rank_quantile

        else:
            if rank_method == 'rank':
                rank = data.rank(pct=True, ascending=positive)
            elif rank_method == 'quantile':
                rank_quantile = (data - data.min()) / (data.max() - data.min())
                if positive:
                    rank = rank_quantile
                else:
                    rank = 1 - rank_quantile

    # rank = np.ceil(rank.astype(float) * scale)
    rank = rank.astype(float) * scale  # 不需要向上取整

    return rank
